/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package perplexity;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 *
 * @author mwh
 */
public class ARPALMEstimator {

    //List<String> file;
    Map<Integer, Map<String, ARPALine>> arpa;


    public ARPALMEstimator(String ARPAfile) throws FileNotFoundException, IOException {

        Integer nowReading = null;
        //      file = new ArrayList<String>();
        BufferedReader BufLM = new BufferedReader(new FileReader(ARPAfile));
        arpa = new HashMap<Integer, Map<String, ARPALine>>();
        //ttern pattern = Pattern.compile("<W.*msd=\"(\\S*)\".*>(.*)<");
        Pattern patternNewNGRAM = Pattern.compile("\\\\(\\d+)-grams:");
        Pattern patternLineWithWeigth = Pattern.compile("(.+)\\t(.+)\\t(.+)");
        Pattern patternLineWithoutWeigth = Pattern.compile("(.+)\\t(.+)");
        System.err.println(patternLineWithWeigth);
        System.err.println(patternNewNGRAM);
        String line;
        while ((line = BufLM.readLine()) != null) {
//            file.add(line);
            //System.err.println(line);
            Matcher ma = patternNewNGRAM.matcher(line);
            if (ma.matches()) {
                nowReading = Integer.valueOf(ma.group(1));
                arpa.put(nowReading, new HashMap<String, ARPALine>());
            }
            ma = patternLineWithWeigth.matcher(line);
            if (ma.matches()) {
                String ngram = ma.group(2);
                double weight = Double.valueOf(ma.group(3));
                double prob = Double.valueOf(ma.group(1));
                ARPALine a = new ARPALine("", prob, weight);
                arpa.get(nowReading).put(ngram, a);
            } else {
                ma = patternLineWithoutWeigth.matcher(line);
                if (ma.matches()) {
                    String ngram = ma.group(2);
                    double weight = 0;
                    double prob = Double.valueOf(ma.group(1));
                    ARPALine a = new ARPALine("", prob, weight);
                    arpa.get(nowReading).put(ngram, a);
                }
            }
        }
    }


    private ARPALine findLine(List<String> tokens) {


//System.err.println("fl: " + tokens);
        StringBuffer strBuf = new StringBuffer();
        for (String token : tokens) {
            strBuf.append(token + " ");
        }
        strBuf.deleteCharAt(strBuf.length() - 1);
        String ngram = strBuf.toString();

        return arpa.get(tokens.size()).get(ngram);
//        
//        for (String line : file) {
//            String[] fields = line.split("\\t");
//            if (fields.length >= 2) {
//                if (fields[1].equals(ngram)) {
//                    if (fields.length > 2) {
//                        return new ARPALine(fields[1], Double.valueOf(fields[0]), Double.valueOf(fields[2]));
//                    }
//                    // No backoff-value means that the value is 1, which is 0 in log.
//                    return new ARPALine(fields[1], Double.valueOf(fields[0]), 0);
//                }
//            }
//        }
//        return null;
    }


    public double getEstimate(List<String> tokens) throws IOException {

     //   System.err.println("\nTrying: " + tokens);
        // Check if the n-gram is present.
        ARPALine a = findLine(tokens);
        if (a != null) {
   //         System.err.println("Found full ngram " + a.logProb);
            return a.logProb;
        }
        if (tokens.size()==1) {
            return Double.NEGATIVE_INFINITY;
        }
        
            List<String> sublist = tokens.subList(0, tokens.size() - 1);
  //          System.err.print("backof: " + sublist + "   n-1-gram: ");
            a = findLine(sublist);
            double backoffWeight = 0;
            if (a != null) {
                
                backoffWeight = a.backofWeight;
  //              System.err.println("Backoff-weight: " + backoffWeight);
            }
            sublist = tokens.subList(1, tokens.size());
  //          System.err.println(sublist + "#");

            return backoffWeight + getEstimate(sublist);
//        
//        System.err.println("No backoff luck");
//        List<String> sublist = tokens.subList(tokens.size() - 1, tokens.size());
//        System.err.println("1-gram: " + sublist);
//        a = findLine(sublist);
//        if (a != null) {
//            return a.logProb;
//        }
//        System.err.println("No luck at all");
//        return 0;
    }


    public static String join(String token, String[] strings) {
        StringBuffer sb = new StringBuffer();

        for (int x = 0; x < (strings.length - 1); x++) {
            sb.append(strings[x]);
            sb.append(token);
        }
        sb.append(strings[strings.length - 1]);

        return (sb.toString());
    }
}
